#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander

import tf.transformations as tfs
import numpy as np
import math

from arm_manager import ArmManager
from gripper_manager import GripperManager
from chassis import ChassisManager
from sensors_init import SensorsInit

from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose
from std_msgs.msg import UInt8
from std_msgs.msg import Float64MultiArray
from webots_ros.msg import RecognitionObject


class MoveItFkDemo:
    def __init__(self):
        # 初始化move_group的API
        moveit_commander.roscpp_initialize(sys.argv)
        # 初始化ROS节点
        rospy.init_node('moveit_fk_demo')
       
        self.setonceflag = 1
        self.enable_flag = 0
        self.objects_info = ['none','can','model']  
        self.camera_to_base_R_T = None
        camera_to_chassis_base_R = None
        self.cal_R_T()

        # # 初始化需要使用move group控制的机械臂中的self.arm group
        # self.arm = moveit_commander.MoveGroupCommander('ls')
                
        # # 获取终端link的名称
        # self.end_effector_link = self.arm.get_end_effector_link()
                        
        # # 设置目标位置所使用的参考坐标系
        # reference_frame = 'base_link'
        # self.arm.set_pose_reference_frame(reference_frame)
                
        # # 当运动规划失败后,允许重新规划
        # self.arm.allow_replanning(True)
        
        # # 设置位置(单位：米)和姿态（单位：弧度）的允许误差
        # self.arm.set_goal_position_tolerance(0.001)
        # self.arm.set_goal_orientation_tolerance(0.001)
        
        #arm
        self.cmarm = ArmManager()
        #gripper
        self.cmgripper = GripperManager()
        #chassis
        self.cm = ChassisManager()
        #camera
        self.si = SensorsInit()
        self.si.enable_sensors()

        rospy.Subscriber('/sign', UInt8, self.get_sign_callback)
        rospy.Subscriber('/joint_group_position_controller/command', Float64MultiArray, self.get_target_pose_callback)
        rospy.Subscriber('/robot/kinect_color/recognition_objects', RecognitionObject, self.get_object_info_callback)
        rospy.spin()
        
    def cal_R_T(self):
            camera_to_chassis_base_R = tfs.euler_matrix(-0.261795,-1.5708,0,'sxyz')
            camera_to_chassis_base_T = (np.array([[0,1.08+0.05,0,1]])).T
            for i in range(3):
                camera_to_chassis_base_R[i,3] = camera_to_chassis_base_T[i,0]

            chassis_base_to_arm_R = (tfs.euler_matrix(-1.5708,0,0)).T
            chassis_base_to_arm_T = -np.dot(chassis_base_to_arm_R,(np.array([[0,0.06,0,1]])).T)
            for i in range(3):
                chassis_base_to_arm_R[i,3] = chassis_base_to_arm_T[i,0]

            self.camera_to_base_R_T = np.dot(chassis_base_to_arm_R,camera_to_chassis_base_R)

    def get_sign_callback(self,msg_sign):
        self.enable_flag = msg_sign.data

    def get_target_pose_callback(self,msg_joint_value):
        self.cmarm.set_velocity(msg_joint_value.data,0.2)
        rospy.sleep(5)
        self.cmgripper.set_velocity(0.5,0.15)
        rospy.sleep(5)

    def get_object_info_callback(self,msg_object_info):
        if(self.objects_info[self.enable_flag] == msg_object_info.model):
            if(abs(msg_object_info.position.x+0.02)>0.01):
            # if(0):
                if(self.setonceflag):
                    if(msg_object_info.position.x>0):
                        self.cm.set_velocity(999,1)
                    else:
                        self.cm.set_velocity(-999,1)
                    self.setonceflag = 0
            #   self.cm.set_velocity(msg_object_info.position.x/(2*0.1),0.5)
            else:
                self.cm.set_velocity(999,0)
                rospy.sleep(2)
                ob_to_camera = (np.array([[msg_object_info.position.x,msg_object_info.position.y,msg_object_info.position.z,1]])).T

                ob_to_base = np.dot(self.camera_to_base_R_T,ob_to_camera)
                print("--------------")
                print(ob_to_base)
                self.set_pose(ob_to_base)
    
    def ik(self,y):
        theta0 = 0.7
        l0 = 0.112
        l1 = 0.6035+0.0415
        l2 = 0.1175
        l3 = 0.16
        f1 = 0.9323
        PI = 3.141592
        theta1 = math.asin((y+theta0-f1)/l1) + PI/2
        theta3 = PI/2 - theta1 
        x = math.cos(theta1-PI/2)*l1 + l0 + l2 + l3
        return [theta0,theta1,0,theta3,0],x

    def set_pose(self,taget_pose):

        joint_positions,x = self.ik(taget_pose[2,0])
        # print(joint_positions)
        # self.arm.set_joint_value_target(joint_positions)
        # self.arm.go()
        # rospy.sleep(2)

        self.cm.set_angle(0,1)
        rospy.sleep(1)

        self.cmarm.set_velocity(joint_positions,1)
        rospy.sleep(1)

        print(taget_pose[0,0],x)
        self.cm.set_position(-(taget_pose[0,0] - x),1)

        self.cmgripper.set_velocity(0.8/0.05*(0.05-0.031),0.15)
        rospy.sleep(8)

        joint_positions,x = self.ik(taget_pose[2,0]+0.1)
        self.cmarm.set_velocity(joint_positions,0.2)
        rospy.sleep(2)

        self.enable_flag = 0
        # 关闭并退出moveit
        # moveit_commander.roscpp_shutdown()
        # moveit_commander.os._exit(0)

if __name__ == "__main__":
    try:
        MoveItFkDemo()
    except rospy.ROSInterruptException:
        pass


# y = 0.6

# theta0 = 0.7
# l0 = 0.112
# l1 = 0.6035+0.0415
# l2 = 0.1175
# l3 = 0.14
# f1 = 0.9323
# PI = 3.141592
# theta1 = math.asin((y+theta0-f1)/l1) + PI/2
# theta3 = PI/2 - theta1 

# x = math.cos(theta1-PI/2)*l1 + l2 + l0

# print(theta1)
# print(theta3)
# print(x)